import wave
import math
import os
import numpy as np

def en(wave_data, window_size, head_size, frame_size, file_path):
    if os.path.exists(file_path):
        os.remove(file_path)
    file_write = open(file_path, 'w')
    start_location = head_size
    end_location = head_size + window_size
    result = []
    while(end_location != frame_size):
        temp_en = 0
        for m in range(start_location, end_location):
            temp_en = temp_en + pow(wave_data[m], 2)
        result += [temp_en]
        file_write.writelines(str(temp_en))
        file_write.write('\n')
        end_location = end_location + window_size
        if end_location < frame_size :
            start_location = start_location + window_size
        else:
            end_location = frame_size
    file_write.close()
    return result

def zero(wave_data, window_size, head_size, frame_size, file_path):
    if os.path.exists(file_path):
        os.remove(file_path)
    file_write = open(file_path, 'w')
    start_location = head_size
    end_location = head_size + window_size
    result = []
    while(end_location != frame_size):
        temp_zero = 0
        for m in range(start_location + 1, end_location):
            if (wave_data[m] * wave_data[m-1]) > 0 :
                temp_zero = temp_zero
            else:
                temp_zero = temp_zero + 1
        result += [temp_zero*window_size]
        file_write.writelines(str(temp_zero*window_size))
        file_write.write('\n')
        end_location = end_location + window_size
        if end_location < frame_size :
            start_location = start_location + window_size
        else:
            end_location = frame_size
    file_write.close()
    return result

def main():
    path = os.path.abspath('.') + '/wav_files/'
    energy_path = os.path.abspath('.') + '/energy_files/'
    zero_rate_path = os.path.abspath('.') + '/zero_rate_files/'
    file_list = os.listdir(path)
    for file in file_list:
        file_info = os.path.splitext(file)
        file_name = file_info[0]
        energy_file_path = energy_path + file_name + '_en.txt'
        zero_rate_file_path = zero_rate_path + file_name + '_zero.txt'
        file_path = path + file
        print(file_path)
        f = wave.open(file_path, 'rb')
        nframes = f.getnframes()
        framerate = f.getframerate()
        temp_data = f.readframes(nframes)
        wave_data = np.fromstring(temp_data, dtype=np.int16)
        f.close()
        en_data = en(wave_data, 256, 44, nframes, energy_file_path)
        zero_data = zero(wave_data, 256, 44, nframes, zero_rate_file_path)

        '''
        可直观显示结果的代码，最终代码无需用到
        time = np.arange(0, len(wave_data)) * (1.0 / framerate)
        time2 = np.arange(0, len(zero_data)) * (1.0 / framerate)
        pl.subplot(211)
        pl.plot(time, wave_data)
        pl.ylabel("Amplitude")
        pl.subplot(212)
        pl.plot(time2, zero_data)
        pl.ylabel("ZCR")
        pl.xlabel("time (seconds)")
        pl.show()
        '''

main()